# encoding: utf-8

import os
import torch
from torch.utils import data
from torchvision import transforms
from PIL import Image
from tqdm import tqdm

from models import ImageGenerator
from utils.images import remove_files

transform = transforms.Compose([
    transforms.Resize((600, 600)),
    transforms.ToTensor(),
])
if not os.path.exists("./models"):
    os.mkdir("./models")

# remove imgs and models
remove_files("./models")
remove_files("./imgsdemo")


# dataloader
class DataGen(data.Dataset):
    def __init__(self, a_path, b_path, transform):
        self.data = []
        self.label = []
        self.transform = transform
        self.load_data(a_path, b_path)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        # img_a = self.transform(self.data[index])
        # img_b = self.transform(self.label[index])
        # img_a = self.transform(Image.open(self.data[index]).convert("RGB"))
        # img_b = self.transform(Image.open(self.label[index]).convert("RGB"))
        # return img_a, img_b
        return self.data[index], self.label[index]

    def load_data(self, a_path, b_path):
        img_a = Image.open(a_path).convert("RGB")
        img_b = Image.open(b_path).convert("RGB")
        self.data.append(self.transform(img_a))
        self.label.append(self.transform(img_b))
        # self.data.append(a_path)
        # self.label.append(b_path)


class DataGen2(data.Dataset):
    def __init__(self, file_path: str):
        self.data = []
        self.label = []
        self.transform = transform
        self.load_data(file_path)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self.data[index], self.label[index]

    def load_data(self, file_path):
        labels = os.listdir(os.path.join(file_path, "label"))
        for file_name in tqdm(labels, desc="加载数据集"):
            # print(file_name)  # AP3361_0.jpg
            if os.path.exists(os.path.join(file_path, "origin", file_name)):
                img_a = self.transform(Image.open(os.path.join(file_path, "origin", file_name)))
                self.data.append(img_a)
                img_b = self.transform(Image.open(os.path.join(file_path, "label", file_name)))
                self.label.append(img_b)


class DataGen3(data.Dataset):
    def __init__(self, file_path: str):
        self.data = []
        self.label = []
        self.transform = transform
        self.load_data(file_path)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self.transform(Image.open(self.data[index]).convert("RGB")), self.transform(
            Image.open(self.label[index]).convert("RGB"))

    def load_data(self, file_path):
        labels = os.listdir(os.path.join(file_path, "label"))
        for file_name in labels:
            # print(file_name)  # AP3361_0.jpg
            if os.path.exists(os.path.join(file_path, "origin", file_name)):
                # img_a = self.transform(Image.open(os.path.join(file_path, "origin", file_name)))
                self.data.append(os.path.join(file_path, "origin", file_name))
                # img_b = self.transform(Image.open(os.path.join(file_path, "label", file_name)))
                self.label.append(os.path.join(file_path, "label", file_name))


# datasets = DataGen("./datasdemo/a.png", "./datasdemo/b.png", transform)
datasets = DataGen2("../img2img_unet/imgs")
# datasets = DataGen3("../img2img_unet/imgs")
batch_size = 8
dataloader = data.DataLoader(datasets, batch_size=batch_size, drop_last=True)

epochs = 32

image_generator = ImageGenerator(epochs=epochs, save_model=True)

image_generator.train(dataloader)
